#include "pch.h"
#include "FragmentBuilder.h"
#include "Helpers.h"

#pragma comment(lib, "mfreadwrite")
#pragma comment(lib, "mfplat")
#pragma comment(lib, "mfuuid")

using namespace Microsoft::Media::FragmentedMediaSink;

FragmentBuilder::FragmentBuilder(IFragmentController^ fragmentController, IMFMediaType* spInputMediaType, unsigned int inputStreamIndex)
	: _inputStreamIndex(inputStreamIndex)
	, _inputMediaType(spInputMediaType)
	, _fragmentController(fragmentController)
	, initiated(false)
	, streamIndex(0)
{
	Helpers::ConvertPropertiesToMediaType(fragmentController->EncodingProperties, &_outputMediaType);

	GUID majorType;
	_outputMediaType->GetMajorType(&majorType);
	isVideo = (majorType == MFMediaType_Video);
}

HRESULT FragmentBuilder::InitializeFragment(_In_ unsigned int fragmentIndex)
{
	HRESULT hr = S_OK;

	//hr = MFStartup(MF_VERSION);

	//ComPtr<IMFAttributes> spAttr;
	ComPtr<IMFByteStream> byteStream;

	if (SUCCEEDED(hr))
	{
		//hr = MFCreateAttributes(&spAttr, 1);
	}

	if (SUCCEEDED(hr))
	{
		// unable to use MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS on Surface RT and Suface Pro 2 devices.
		//hr = spAttr->SetUINT32(MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, true);
		//hr = spAttr->SetUINT32(MF_SINK_WRITER_DISABLE_THROTTLING, true);
	}

	if (SUCCEEDED(hr))
	{
		auto stream = _fragmentController->StartFragment();

		hr = MFCreateMFByteStreamOnStreamEx((IUnknown*)stream, &byteStream);

		if (SUCCEEDED(hr))
		{
			fragment = ref new Fragment(stream, _inputStreamIndex, fragmentIndex);
		}
	}

	if (SUCCEEDED(hr))
	{
		hr = MFCreateSinkWriterFromURL(L".mp4", byteStream.Get(), NULL, &sinkWriter);
	}

	if (SUCCEEDED(hr))
	{
		hr = sinkWriter->AddStream(_outputMediaType.Get(), &streamIndex);
	}

	// Set the input media type.
	if (SUCCEEDED(hr))
	{
		hr = sinkWriter->SetInputMediaType(streamIndex, _inputMediaType.Get(), NULL);
	}

	// Tell the sink writer to start accepting data.
	if (SUCCEEDED(hr))
	{
		hr = sinkWriter->BeginWriting();
	}

	if (SUCCEEDED(hr))
	{
		initiated = true;
	}

	return hr;
}

HRESULT FragmentBuilder::TerminateFragment(long long *startTime, long long *endTime)
{
	HRESULT hr = S_OK;

	if (initiated)
	{
		initiated = false;

		hr = sinkWriter->Finalize();

		if (SUCCEEDED(hr))
		{
			fragment->StartTime = *startTime;
			fragment->EndTime = *endTime;
			_fragmentController->FragmentCompleted(fragment);
			*startTime = fragment->StartTime;
			*endTime = fragment->EndTime;
		}

		sinkWriter = nullptr;
	}

	return hr;
}

HRESULT FragmentBuilder::WriteSample(_In_ IMFSample* spSample)
{
	HRESULT hr = S_OK;
	
	if (initiated) // safe guard against samples coming in after EOS.
	{
		hr = sinkWriter->WriteSample(streamIndex, spSample);
	}

	return hr;
}

HRESULT FragmentBuilder::SendTick(_In_ long long timeStamp)
{
	HRESULT hr = S_OK;

	if (initiated)
	{
		hr = sinkWriter->SendStreamTick(streamIndex, timeStamp);
	}

	return hr;
}

HRESULT FragmentBuilder::Flush()
{
	HRESULT hr = S_OK;

	if (initiated)
	{
		hr = sinkWriter->Flush(streamIndex);
	}

	return hr;
}